"""
Combined Improved Dung Beetle Optimization (IDBO) and Extreme Learning Machine (ELM) framework
for precise SOC estimation of batteries

Implementation based on the paper:
"A combined improved dung beetle optimization and extreme learning machine framework for precise SOC estimation"
"""

import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
import pandas as pd
from typing import Tuple, List, Dict, Any, Optional

from improved_dbo import ImprovedDungBeetleOptimizer
from elm import ExtremeLearningMachine


class IDBO_ELM_SOC:
    def __init__(
        self,
        hidden_neurons: int = 20,
        activation: str = 'sigmoid',
        population_size: int = 30,
        max_iterations: int = 100,
        initial_rolling_factor: float = 0.5,
        communication_probability: float = 0.3,
        elite_group_size: int = 5,
        opposition_probability: float = 0.1,
        levy_alpha: float = 1.5,
        golden_ratio: float = 1.618,
        random_state: Optional[int] = None
    ):
        """
        Initialize the IDBO-ELM framework for SOC estimation
        
        Args:
            hidden_neurons: Number of hidden neurons in ELM
            activation: Activation function for ELM
            population_size: Population size for IDBO
            max_iterations: Maximum iterations for IDBO
            initial_rolling_factor: Initial factor for rolling behavior
            communication_probability: Probability of communication
            elite_group_size: Size of elite group in IDBO
            opposition_probability: Probability of opposition-based learning
            levy_alpha: Parameter for Lévy flight
            golden_ratio: Golden ratio for sine strategy
            random_state: Random seed for reproducibility
        """
        self.hidden_neurons = hidden_neurons
        self.activation = activation
        self.population_size = population_size
        self.max_iterations = max_iterations
        self.initial_rolling_factor = initial_rolling_factor
        self.communication_probability = communication_probability
        self.elite_group_size = elite_group_size
        self.opposition_probability = opposition_probability
        self.levy_alpha = levy_alpha
        self.golden_ratio = golden_ratio
        self.random_state = random_state
        
        # These will be initialized later when data dimensions are known
        self.elm = None
        self.optimizer = None
        self.input_size = None
        self.output_size = None
        
        # For storing results
        self.best_params = None
        self.best_fitness = None
        self.convergence_curve = None
        
        # Scalers for data normalization (0-1 range as mentioned in the paper)
        self.X_scaler = MinMaxScaler(feature_range=(0, 1))
        self.y_scaler = MinMaxScaler(feature_range=(0, 1))
    
    def _objective_function(self, params: np.ndarray) -> float:
        """
        Objective function for IDBO optimization
        Evaluates ELM performance with given parameters (input weights and biases)
        
        Args:
            params: Flattened array of ELM parameters (input weights and biases)
            
        Returns:
            Error metric (combined RMSE, MAE, R²) to be minimized
        """
        # Set ELM parameters (input weights and biases)
        self.elm.set_params(params)
        
        # Train ELM with current parameters (this calculates output weights β analytically)
        self.elm.fit(self.X_train_scaled, self.y_train_scaled)
        
        # Predict on validation set
        y_pred = self.elm.predict(self.X_val_scaled)
        
        # Calculate error metrics
        rmse, mae, r2 = self._calculate_metrics(y_pred, self.y_val_scaled)
        
        # Combined fitness function (optimized weights)
        # Lower RMSE and MAE are better, higher R² is better
        fitness = 0.6 * rmse + 0.3 * mae - 0.5 * r2
        
        return fitness
    
    def _calculate_metrics(self, y_pred: np.ndarray, y_true: np.ndarray) -> Tuple[float, float, float]:
        """
        Calculate evaluation metrics
        
        Args:
            y_pred: Predicted values
            y_true: True values
            
        Returns:
            Tuple of (RMSE, MAE, R²)
        """
        # Root Mean Squared Error
        rmse = np.sqrt(np.mean((y_pred - y_true) ** 2))
        
        # Mean Absolute Error
        mae = np.mean(np.abs(y_pred - y_true))
        
        # R² score with protection against division by zero
        ss_total = np.sum((y_true - np.mean(y_true)) ** 2)
        ss_residual = np.sum((y_true - y_pred) ** 2)
        
        # Check if ss_total is close to zero (constant target values)
        if ss_total < 1e-10:
            # If target values are constant, R² is undefined
            # Return 0 or 1 based on prediction accuracy
            if ss_residual < 1e-10:
                r2 = 1.0  # Perfect prediction
            else:
                r2 = 0.0  # Imperfect prediction
        else:
            r2 = 1 - (ss_residual / ss_total)
            # Clip R² to valid range [-1, 1]
            r2 = np.clip(r2, -1.0, 1.0)
        
        return rmse, mae, r2
    
    def preprocess_data(self, X: np.ndarray, y: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """
        Preprocess the data for training and testing
        Normalizes data to [0,1] range as mentioned in the paper
        
        Args:
            X: Input features
            y: Target values (SOC)
            
        Returns:
            Tuple of (X_train, X_val, X_test, y_train, y_val, y_test)
        """
        # Check for NaN values and replace them
        if np.isnan(X).any():
            print(f"Warning: Found {np.isnan(X).sum()} NaN values in features. Replacing with zeros.")
            X = np.nan_to_num(X, nan=0.0)
        
        if np.isnan(y).any():
            print(f"Warning: Found {np.isnan(y).sum()} NaN values in target. Replacing with zeros.")
            y = np.nan_to_num(y, nan=0.0)
        
        # Check for infinite values and replace them
        if np.isinf(X).any():
            print(f"Warning: Found {np.isinf(X).sum()} infinite values in features. Replacing with large values.")
            X = np.nan_to_num(X, posinf=1e10, neginf=-1e10)
        
        if np.isinf(y).any():
            print(f"Warning: Found {np.isinf(y).sum()} infinite values in target. Replacing with boundary values.")
            y = np.nan_to_num(y, posinf=100.0, neginf=0.0)
        
        # Scale input features and target values to [0,1] range
        X_scaled = self.X_scaler.fit_transform(X)
        y_scaled = self.y_scaler.fit_transform(y.reshape(-1, 1))
        
        # Split data into training and testing sets (80% train, 20% test)
        X_train, X_test, y_train, y_test = train_test_split(
            X_scaled, y_scaled, test_size=0.2, random_state=self.random_state
        )
        
        # Further split training data into training and validation sets (80% train, 20% validation)
        X_train, X_val, y_train, y_val = train_test_split(
            X_train, y_train, test_size=0.2, random_state=self.random_state
        )
        
        return X_train, X_val, X_test, y_train, y_val, y_test
    
    def fit(self, X: np.ndarray, y: np.ndarray) -> Dict[str, Any]:
        """
        Train the IDBO-ELM framework following the paper's algorithm
        
        Args:
            X: Input features
            y: Target values (SOC)
            
        Returns:
            Dictionary with training results
        """
        # Step 1: Get data dimensions
        self.input_size = X.shape[1]
        self.output_size = 1  # SOC is a single value
        
        # Step 2: Preprocess data (normalize to [0,1] range)
        self.X_train_scaled, self.X_val_scaled, self.X_test_scaled, \
        self.y_train_scaled, self.y_val_scaled, self.y_test_scaled = self.preprocess_data(X, y)
        
        print(f"Data preprocessed. Training samples: {self.X_train_scaled.shape[0]}, "
              f"Validation samples: {self.X_val_scaled.shape[0]}, "
              f"Test samples: {self.X_test_scaled.shape[0]}")
        
        # Step 3: Initialize ELM with random input weights and biases
        self.elm = ExtremeLearningMachine(
            input_size=self.input_size,
            hidden_size=self.hidden_neurons,
            output_size=self.output_size,
            activation=self.activation,
            random_state=self.random_state
        )
        
        # Calculate parameter dimensions for IDBO
        # Parameters to optimize: input weights (input_size * hidden_neurons) and biases (hidden_neurons)
        param_size = self.input_size * self.hidden_neurons + self.hidden_neurons
        print(f"Parameter space dimension: {param_size}")
        
        # Step 4: Initialize IDBO optimizer with the improved features
        self.optimizer = ImprovedDungBeetleOptimizer(
            objective_function=self._objective_function,
            dimensions=param_size,
            population_size=self.population_size,
            max_iterations=self.max_iterations,
            lower_bound=-1.0,  # Normalized parameter range
            upper_bound=1.0,
            initial_rolling_factor=self.initial_rolling_factor,
            communication_probability=self.communication_probability,
            elite_group_size=self.elite_group_size,
            opposition_probability=self.opposition_probability,
            levy_alpha=self.levy_alpha,
            golden_ratio=self.golden_ratio
        )
        
        # Step 5: Run IDBO optimization to find optimal input weights and biases
        print("Starting IDBO-ELM optimization for SOC estimation...")
        self.best_params, self.best_fitness, self.convergence_curve = self.optimizer.optimize()
        
        print(f"Optimization completed. Best fitness: {self.best_fitness}")
        
        # Step 6: Set the best parameters to ELM and train final model
        self.elm.set_params(self.best_params)
        
        # Train ELM with optimal parameters to calculate output weights β
        self.elm.fit(self.X_train_scaled, self.y_train_scaled)
        
        # Evaluate on test set
        test_rmse, test_mae, test_r2 = self.evaluate(self.X_test_scaled, self.y_test_scaled)
        
        print(f"Final model evaluation on test set:")
        print(f"RMSE: {test_rmse:.6f}")
        print(f"MAE: {test_mae:.6f}")
        print(f"R²: {test_r2:.6f}")
        
        # Return results
        results = {
            'best_fitness': self.best_fitness,
            'convergence_curve': self.convergence_curve,
            'test_rmse': test_rmse,
            'test_mae': test_mae,
            'test_r2': test_r2
        }
        
        return results
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Predict SOC values for new data
        
        Args:
            X: Input features
            
        Returns:
            Predicted SOC values
        """
        # Check for NaN values and replace them
        if np.isnan(X).any():
            X = np.nan_to_num(X, nan=0.0)
        
        # Scale input
        X_scaled = self.X_scaler.transform(X)
        
        # Make predictions
        y_pred_scaled = self.elm.predict(X_scaled)
        
        # Inverse transform to original scale
        y_pred = self.y_scaler.inverse_transform(y_pred_scaled)
        
        return y_pred
    
    def evaluate(self, X: np.ndarray, y: np.ndarray) -> Tuple[float, float, float]:
        """
        Evaluate the model performance
        
        Args:
            X: Test data (already scaled)
            y: True values (already scaled)
            
        Returns:
            Tuple of (RMSE, MAE, R²)
        """
        # Predict
        y_pred = self.elm.predict(X)
        
        # Calculate metrics
        return self._calculate_metrics(y_pred, y)
    
    def plot_convergence(self) -> None:
        """
        Plot the convergence curve of the optimization process
        """
        plt.figure(figsize=(10, 6))
        plt.plot(self.convergence_curve)
        plt.title('IDBO-ELM Convergence Curve')
        plt.xlabel('Iteration')
        plt.ylabel('Fitness Value')
        plt.grid(True)
        plt.savefig('idbo_elm_convergence.png')
        plt.show()
    
    def plot_predictions(self, X_test: np.ndarray, y_test: np.ndarray) -> None:
        """
        Plot the predicted vs actual SOC values
        
        Args:
            X_test: Test input features (original scale)
            y_test: True SOC values (original scale)
        """
        try:
            # Make predictions
            y_pred = self.predict(X_test)
            
            # Create time index for plotting
            time_index = np.arange(len(y_test))
            
            # Plot predictions vs actual
            plt.figure(figsize=(12, 6))
            plt.plot(time_index, y_test, 'b-', label='Actual SOC')
            plt.plot(time_index, y_pred, 'r--', label='Predicted SOC')
            plt.title('SOC Estimation: Actual vs Predicted')
            plt.xlabel('Time Step')
            plt.ylabel('SOC (%)')
            plt.legend()
            plt.grid(True)
            plt.savefig('soc_prediction.png')
            
            # Plot error
            plt.figure(figsize=(12, 6))
            plt.plot(time_index, np.abs(y_test.flatten() - y_pred.flatten()), 'g-')
            plt.title('Absolute SOC Estimation Error')
            plt.xlabel('Time Step')
            plt.ylabel('Absolute Error (%)')
            plt.grid(True)
            plt.savefig('soc_error.png')
            
            # Plot scatter
            plt.figure(figsize=(8, 8))
            plt.scatter(y_test, y_pred, alpha=0.5)
            plt.plot([0, 100], [0, 100], 'r--')  # Perfect prediction line
            plt.title('Actual vs Predicted SOC')
            plt.xlabel('Actual SOC (%)')
            plt.ylabel('Predicted SOC (%)')
            plt.grid(True)
            plt.axis('equal')
            plt.savefig('soc_scatter.png')
            
            plt.show()
        except Exception as e:
            print(f"Error plotting predictions: {e}")
            import traceback
            traceback.print_exc()